/*
  Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

#include "analytic_case.h"
#include "blackbody.h"
#include "grain_model.h"
#include "misc.h"

void add_cam(T_float theta, T_float phi, T_float cd, T_float fov_fact, 
	     T_float grid_extent,
	     std::vector<vec3d>& pos, 
	     std::vector<vec3d>& dir,
	     std::vector<vec3d>& up, 
	     std::vector<T_float>& fov)
{
  dir.push_back(-vec3d(sin(theta*constants::pi/180.)*cos(phi),
		       sin(theta*constants::pi/180)*sin(phi),
		       cos(theta*constants::pi/180)));
  pos.push_back(-dir.back()*cd);
  up.push_back(vec3d(0,0,1));
  fov.push_back(fov_fact*sqrt(2)*grid_extent/cd);
}

/** This factory draws the random high or low-density cells in a cubic
    distribution. It's just used by the wg96_factory as it makes it
    easier to refine the cells. */  
class clumpy_factory {
public:
  typedef T_grid::T_grid_impl::T_data T_data;
  typedef refinement_accuracy_data<T_data> T_racc;
  typedef cell_tracker<T_data> T_cell_tracker;

  const T_float ff, contrast;
  const T_densities rho_h, rho_1, rho_2, rho0;
  const int maxlevel;

  clumpy_factory (T_float f, T_float c,
		  const T_densities& rho, int ml):
    ff (f), contrast (c), rho_h (rho),
    rho_1 (rho_h/((1-ff)+ff*contrast)), 
    rho_2 (rho_1*contrast), rho0 (rho_1*0.0), maxlevel(ml) {
    cout << "Densities are " << rho_1 << " and " << rho_2 << endl;
  };
  
  bool refine_cell_p (const T_cell_tracker& c) {
    return c.code().level()<maxlevel;};
  bool unify_cell_p (const T_cell_tracker& c, const T_racc& racc) {
    return false;};
  T_data get_data (const T_cell_tracker& c) {
    // center is always low-density to avoid blacking out everything
    return T_data (T_data::T_emitter(),
		   T_data::T_absorber
		   ( ( (mcrx::rnd () > ff) || 
		       (mag(c.getcenter()/c.getsize())<1) ) ? rho_1 : rho_2,
		     vec3d(0,0,0), T_lambda()));
  }
    
   int n_threads () {return 1;};
};


class wg96_factory {
public:
  typedef T_grid::T_grid_impl::T_data T_data;
  typedef refinement_accuracy_data<T_data> T_racc;
  typedef cell_tracker<T_data> T_cell_tracker;

  const T_float radius;
  const int max_level, clump_level;
  const int n_lambda;
  const int n_threads_;
  const T_densities n_tol_;

  clumpy_factory factory;
  adaptive_grid<T_grid::T_grid_impl::T_data> g;

  wg96_factory (T_float f, T_float c, T_float r, 
		const T_densities& rho, 
		const T_densities& ntol, int cl, int ml, int nl, int nt):
    radius (r), max_level (ml), clump_level(cl), n_lambda(nl), n_threads_(nt),
    factory (f, c, rho, cl), n_tol_(ntol),
    g (vec3d (-r,-r,-r ), vec3d (r,r,r ), factory)
  {};
  
  const T_densities& get_rho(const vec3d& pos) {
    if (mag(pos) > radius) 
      return factory.rho0;

    const T_grid::T_cell_tracker gc = g.locate(pos, 0, false);
    if(gc)
      return gc.data()->get_absorber().densities();
    else
      return factory.rho0;
  };

  bool refine_cell_p (const T_cell_tracker& c) {
    const int level = c.code().level();
    const vec3d center = c.getcenter();
    const T_float r = sqrt(dot (center, center));
    const T_densities& rho(get_rho(center));
    T_densities n2 (dot(c.getsize(), c.getsize())*rho*rho);

    // refine the cell if it overlaps with spherical surface or if it
    // exceeds our n_tol

    // see if the cell possibly overlaps with the spherical surface or
    // if the optical depth exceeds tolerance since only in that case
    // it needs to be refined beyond the clump level
    if( (level<= clump_level) ||
	( (level < max_level) &&
	  ( (fabs(radius-r) <= sqrt(3)/2*c.getsize()[0]) ||
	    all(n2>(n_tol_*n_tol_)) ) ) )
      return true;
    else return false;
  };

  bool unify_cell_p (const T_cell_tracker& c, const T_racc& racc) {
    return false;
  };

  T_data get_data (const T_cell_tracker& c) {
    const vec3d center = c.getcenter();
    return T_data (T_data::T_emitter(), 
		   T_data::T_absorber(get_rho(center),
				      vec3d(0,0,0), 
				      T_lambda(n_lambda)));
  }
    
   int n_threads () {return n_threads_;};
};


void setup (po::variables_map& opt)
{
  T_unit_map units;
  units ["length"] = "au";
  units ["mass"] = "Msun";
  units ["luminosity"]= "W";
  units ["wavelength"] = "m";
  units ["L_lambda"] = "W/m";

  const int gridlevel = opt["gridlevel"].as<int>();
  const int maxlevel = opt["maxlevel"].as<int>();
  const vec3d grid_extent(1000,1000,1000);
  const T_float ff=opt["ff"].as<T_float>();
  const T_float contrast=opt["contrast"].as<T_float>();

  const T_float source_temp = opt["source_temp"].as<T_float>();
  const T_float source_lum = opt["source_lum"].as<T_float>();
  const int npix = opt["npix"].as<int>();
  const int n_threads = opt["n_threads"].as<int>();
  const T_float tau = opt["tau"].as<T_float>();
  const T_float tau_tol = opt["tau_tol"].as<T_float>();
  const T_float cameradist = 50000;
  const string grain_model = opt["grain_model"].as<string>();

  // we need a preferences object to pass parameters to grain model
  Preferences p;
  p.setValue("grain_data_directory", 
	     word_expand(opt["grain_data_dir"].as<string>())[0]);
  p.setValue("grain_model", grain_model);
  p.setValue("wd01_parameter_set", opt["wd01_parameter_set"].as<string>());
  p.setValue("template_pah_fraction", 0.5);
  p.setValue("use_dl07_opacities", true);
  p.setValue("n_threads", n_threads);
  p.setValue("use_grain_temp_lookup", true);
  p.setValue("use_cuda", false);

  // read wavelengths
  array_1 lambda;
  lambda.reference(logspace(opt["lambda_min"].as<T_float>(),
			    opt["lambda_max"].as<T_float>(),
			    opt["n_lambda"].as<int>()));
  // lambdas.push_back(0.4e-6);
  // lambdas.push_back(0.55e-6);
  // lambdas.push_back(0.7e-6);

  // load grain model
  T_dust_model::T_scatterer_vector sv;
  sv.push_back(load_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(p,units));
  T_dust_model model (sv);
  model.set_wavelength(lambda);

  std::vector<T_float> lambdas(lambda.begin(), lambda.end());
  const int lambda550 = lower_bound(lambdas.begin(), lambdas.end(),
				    0.551e-6) - lambdas.begin()-1;

  // find opacity normalization for whatever dust we're using
  const T_float unit_opacity_rho = 1.0/sv[0]->opacity()(lambda550);
  T_densities rho(1); rho=unit_opacity_rho/grid_extent*tau;
  T_densities n_tol(1); n_tol = tau_tol*unit_opacity_rho;

  // set seed for factory
  mcrx::seed(opt["grid_seed"].as<int>());
  wg96_factory factory (ff, contrast, grid_extent[0], rho, n_tol,
			gridlevel, maxlevel, lambda.size(), n_threads);

  blackbody bb(source_temp, 3.91e26*source_lum);
  boost::shared_ptr<T_emission> em
    (new pointsource_emission<polychromatic_policy, T_rng_policy>
     (vec3d (.000, .00, -.00), bb.emission(lambda)));

  // define cameras
  std::vector<vec3d> pos,dir,up;
  std::vector<T_float> fov;
  add_cam(77.5, 18.233, cameradist, 1.5, grid_extent[0], pos, dir, up, fov);
  //add_cam(82.5, 18.233, cameradist, 1.5, grid_extent[0], pos, dir, up, fov);
  //add_cam(87.5, 18.233, cameradist, 1.5, grid_extent[0], pos, dir, up, fov);
  boost::shared_ptr<T_emergence> cameras
    (new T_emergence(pos,dir,up,fov, npix));

  run_case(opt, 
	   factory,
	   units,
	   lambda,
	   -grid_extent,
	   grid_extent,
	   *em,
	   *cameras,
	   model); 
}


void pre_shoot(T_xfer&) {};

std::string description()
{
  return "Witt & Gordon (96) clumpy region setup.";
}

void add_custom_options(po::options_description& desc)
{
  desc.add_options()
    ("tau", po::value<T_float>()->default_value(160),
     "equivalent homogeneous V-band optical depth of cloud")
    ("ff", po::value<T_float>()->default_value(0.2),
     "volume filling factor of high-density cells")
    ("contrast", po::value<T_float>()->default_value(800),
     "density contrast between high- and low-density cells")
    ("gridlevel", po::value<int>()->default_value(5), 
     "refinement level of high/low cells")
    ("grid_seed", po::value<int>()->default_value(42), "random number seed for grid generation")
    ("source_temp", po::value<T_float>()->default_value(5800),
     "temperature of central blackbody source")
    ("source_lum", po::value<T_float>()->default_value(1),
     "luminosity of central source (in Lsun)")
    ("grain_model", po::value<string>()->default_value("wd01_Brent_PAH"), "dust grain model (use \"blackbody\" for blackbody grains)")
    ;
}

void print_custom_options(const po::variables_map& opt)
{
  std::cout
       << "tau = " << opt["tau"].as<T_float>() << '\n'
       << "ff = " << opt["ff"].as<T_float>() << '\n'
       << "contrast = " << opt["contrast"].as<T_float>() << '\n'
       << "gridlevel = " << opt["gridlevel"].as<int>() << '\n'
       << "grid_seed = " << opt["grid_seed"].as<int>() << '\n'
       << "source_temp = " << opt["source_temp"].as<T_float>() << '\n'
       << "source_lum = " << opt["source_lum"].as<T_float>() << '\n'
       << "grain_model = " << opt["grain_model"].as<std::string>() << '\n'
    ;
}
